package com.cloud.job.doris;

import cn.hutool.core.lang.UUID;
import com.baomidou.mybatisplus.annotation.TableField;
import com.baomidou.mybatisplus.annotation.TableName;
import com.cloud.core.exception.MyException;
import com.cloud.job.config.DorisConfig;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.codec.binary.Base64;
import org.apache.http.HttpHeaders;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPut;
import org.apache.http.entity.ContentType;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.DefaultRedirectStrategy;
import org.apache.http.impl.client.HttpClientBuilder;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.util.EntityUtils;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.lang.reflect.Field;
import java.nio.charset.StandardCharsets;

/**
 * @author 杨斌
 * 2024-09-13/9/2024
 */
@Component
@Slf4j
public class DorisStreamLoad {
    @Resource
    private DorisConfig dorisConfig;


    private String getLoadUrl(Class<?> clazz) {
        String tableName = getTableName(clazz);
        return String.format("http://%s:%s/api/%s/%s/_stream_load",
                dorisConfig.getHost(),
                dorisConfig.getPort(),
                dorisConfig.getDb(),
                tableName);
    }

    private String getTableName(Class<?> clazz) {
        TableName tableNameAnnotation = clazz.getAnnotation(TableName.class);
        if (tableNameAnnotation != null) {
            return tableNameAnnotation.value();
        }
        throw new MyException("Class " + clazz.getName() + " does not have a @TableName annotation.");
    }


    private String basicAuthHeader(String username, String password) {
        final String tobeEncode = username + ":" + password;
        byte[] encoded = Base64.encodeBase64(tobeEncode.getBytes(StandardCharsets.UTF_8));
        return "Basic " + new String(encoded);
    }

    //构建HTTP客户端
    private final static HttpClientBuilder httpClientBuilder = HttpClients
            .custom()
            .setRedirectStrategy(new DefaultRedirectStrategy() {
                @Override
                protected boolean isRedirectable(String method) {
                    // If the connection target is FE, you need to deal with 307 redirect。
                    return true;
                }
            });

    /**
     * JSON格式的数据导入
     * @param jsonData
     * @throws Exception
     */
    public void loadJson(String jsonData, Class<?> clazz) throws Exception {
        try (CloseableHttpClient client = httpClientBuilder.build()) {
            HttpPut put = new HttpPut(getLoadUrl(clazz));
            put.removeHeaders(HttpHeaders.CONTENT_LENGTH);
            put.removeHeaders(HttpHeaders.TRANSFER_ENCODING);
            put.setHeader(HttpHeaders.EXPECT, "100-continue");
            put.setHeader(HttpHeaders.AUTHORIZATION, basicAuthHeader(dorisConfig.getUserName(), dorisConfig.getPassword()));

            // 设置流加载相关的属性
            put.setHeader("label", UUID.randomUUID().toString());
            put.setHeader("column_separator", ",");
            put.setHeader("format", "json");
            put.setHeader("strip_outer_array", "true");
            put.setHeader("read_json_by_line", "true");

            // 动态生成 jsonpaths 和 columns
            String jsonPaths = generateJsonPaths(clazz);
            String columns = generateColumns(clazz);

            put.setHeader("jsonpaths", jsonPaths);
            put.setHeader("columns", columns);

            // 设置导入
            StringEntity entity = new StringEntity(jsonData, ContentType.create("application/json", StandardCharsets.UTF_8));
            put.setEntity(entity);

            try (CloseableHttpResponse response = client.execute(put)) {
                String loadResult = "";
                if (response.getEntity() != null) {
                    loadResult = EntityUtils.toString(response.getEntity());
                }
                final int statusCode = response.getStatusLine().getStatusCode();
                if (statusCode != 200) {
                    throw new IOException(String.format("Stream load failed. status: %s load result: %s", statusCode, loadResult));
                }
                log.info("Get load result: " + loadResult);
            }catch (Exception e) {
                log.error("Doris流加载发生异常: ", e);
            }
        }
    }

    private String generateJsonPaths(Class<?> clazz) {
        Field[] fields = clazz.getDeclaredFields();
        StringBuilder jsonPaths = new StringBuilder("[");
        for (int i = 0; i < fields.length; i++) {
            if (i > 0) jsonPaths.append(",");
            jsonPaths.append("\"$.").append(fields[i].getName()).append("\"");
        }
        jsonPaths.append("]");
        return jsonPaths.toString();
    }

    private String generateColumns(Class<?> clazz) {
        Field[] fields = clazz.getDeclaredFields();
        StringBuilder columns = new StringBuilder();
        for (int i = 0; i < fields.length; i++) {
            if (i > 0) columns.append(",");
            TableField tableField = fields[i].getAnnotation(TableField.class);
            if (tableField != null && !tableField.value().isEmpty()) {
                columns.append(tableField.value());
            } else {
                columns.append(fields[i].getName());
            }
        }
        return columns.toString();
    }
}

